#ifdef USE_GPU
#ifndef CUDA_COMMON_HPP
#define CUDA_COMMON_HPP

/****************** CUDA KERNEL 声明 ******************/
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>

/// @brief 内存拷贝, cudaMemcpy的封装
/// @tparam idx_t 索引类型
/// @tparam data_t 数据类型
/// @param dst 目的地址
/// @param src 源地址
/// @param n 拷贝的元素个数
/// @param direction 0: Host to Device, 1: Device to Host, 2: Device to Device
template<typename idx_t, typename data_t>
void cuda_memcpy(data_t * dst, const data_t * src, idx_t n, int direction);

/// @brief 分配内存, cudaMalloc的封装
/// @tparam idx_t 索引类型
/// @param data 待分配的内存地址的指针
/// @param n_bytes 分配的字节数
template<typename idx_t>
void cuda_malloc(void **data, idx_t n_bytes);

/// @brief 内存初始化, cudaMemset的封装
/// @tparam idx_t 索引类型
/// @param data 待初始化的内存地址
/// @param val 初始化的值
/// @param n_bytes 初始化的字节数
template<typename idx_t>
void cuda_memset(void *data, int val, idx_t n_bytes);

/// @brief 释放内存, cudaFree的封装
/// @param data 待释放的内存地址
void cuda_free(void *data);

/// @brief 向量内积 res = x dot y
/// @tparam idx_t 索引类型
/// @tparam data_t 数据类型
/// @tparam res_t 点积结果类型，同时也是计算类型
/// @param x_data 向量x的数据
/// @param y_data 向量y的数据
/// @param n 向量的长度
/// @return 点积结果
template<typename idx_t, typename data_t, typename res_t>
res_t cuda_vec_dot(const data_t * x_data, const data_t * y_data, const idx_t n);

/// @brief 向量加法 dst = src1 + src2 * alpha
/// @tparam idx_t 索引类型
/// @tparam data_t 数据类型
/// @tparam scalar_t 系数类型
/// @param dst 目的向量
/// @param src1 源向量1
/// @param src2 源向量2
/// @param alpha 系数
/// @param n 向量长度
template<typename idx_t, typename data_t, typename scalar_t>
void cuda_vec_add(data_t * dst, const data_t * src1, const data_t * src2, const scalar_t alpha, const idx_t n);

/// @brief 向量缩放 dst = src * coeff
/// @tparam idx_t 索引类型
/// @tparam data_t 数据类型
/// @tparam scalar_t 系数类型
/// @param dst 目的向量
/// @param src 源向量
/// @param coeff 系数
/// @param n 向量长度
template<typename idx_t, typename data_t, typename scalar_t>
void cuda_vec_scale(data_t * dst, const data_t * src, const scalar_t coeff, const idx_t n);

/// @brief 向量拷贝 dst = src
/// @tparam idx_t 索引类型
/// @tparam data_t1 目的向量数据类型
/// @tparam data_t2 源向量数据类型
/// @param dst 目的向量
/// @param src 源向量
/// @param n 向量长度
template<typename idx_t, typename data_t1, typename data_t2>
void cuda_vec_copy(data_t1 * dst, const data_t2 * src, const idx_t n);

/// @brief 向量按位乘法 dst[i] = src[i] * coeff[i]
/// @tparam idx_t 索引类型
/// @tparam data_t 数据类型
/// @tparam calc_t 计算类型
/// @param dst 目的向量
/// @param src 源向量
/// @param coeff 系数向量
/// @param n 向量长度
template<typename idx_t, typename data_t, typename calc_t>
void cuda_vec_elem_mul(calc_t * dst, const calc_t * src, const data_t * coeff, const idx_t n);

/// @brief 从 CSR 提取对角线元素到 diag, 用于 Jacobi 的 GPU 加速
/// @tparam idx_t 索引类型
/// @tparam data_t 数据类型
/// @tparam setup_t 用于 Setup 的矩阵元素的数据类型
/// @param diag 对角线元素
/// @param row_ptr 行指针
/// @param col_idx 列索引
/// @param vals 矩阵元素
/// @param n 矩阵行数
template<typename idx_t, typename data_t, typename setup_t>
void cuda_separate_invD(data_t * diag, const idx_t * row_ptr, const idx_t * col_idx, const setup_t * vals, const idx_t n);

/// @brief 稀疏矩阵向量乘法 y = alpha * A * x + beta * y
/// @tparam idx_t 索引类型
/// @tparam data_t 矩阵元素数据类型
/// @tparam calc_t 计算类型(向量数据类型)
/// @param row_ptr 行指针
/// @param col_idx 列索引
/// @param vals 矩阵元素
/// @param x, b, y 向量
/// @param alpha, beta 系数
/// @param n 向量长度/矩阵行数
template<typename idx_t, typename data_t, typename calc_t>
void cuda_spmv(const idx_t * row_ptr, const idx_t * col_idx, const data_t * vals, const calc_t * x, const calc_t * b, calc_t * y, const calc_t alpha, const calc_t beta, const idx_t n);

/// @brief CG 算法 kernel fusion, 减少 kernel launch 开销
///        详见 iter_solver.hpp
/// @tparam idx_t 索引类型
/// @tparam data_t 矩阵元素数据类型
/// @tparam calc_t 计算类型(向量数据类型)
/// @tparam ksp_t Krylov子空间算法中的数据类型
/// @param x, p, r, s, u CG算法中的辅助向量
/// @param alpha, beta CG算法中的辅助参数
/// @param diag Jacobi保存的矩阵对角线元素
/// @param tmp_glb 记录向量点积, tmp_glb[0] = r dot u, tmp_glb[1] = r dot r
/// @param n 向量长度/矩阵行数
template<typename idx_t, typename data_t, typename calc_t, typename ksp_t>
void CG_kernel_1(calc_t * x, const calc_t alpha, calc_t * p, calc_t * r, calc_t * s, calc_t * u, const data_t * diag, ksp_t * tmp_glb, const idx_t n);

// @brief CG 算法 kernel fusion, 减少 kernel launch 开销
///        详见 iter_solver.hpp
/// @tparam idx_t 索引类型
/// @tparam calc_t 计算类型(向量数据类型)
/// @tparam ksp_t Krylov子空间算法中的数据类型
/// @param u, r, s, p, b CG算法中的向量
/// @param tmp_glb 记录向量点积, tmp_glb[0] = r dot u, tmp_glb[1] = s dot p
///                             tmp_glb[2] = b dot b, tmp_glb[3] = r dot r
/// @param n 向量长度/矩阵行数
template<typename idx_t, typename calc_t, typename ksp_t>
void CG_kernel_2(calc_t * u, calc_t * r, calc_t * s, calc_t * p, const calc_t * b, ksp_t * tmp_glb, const idx_t n);

#endif
#endif